#!/usr/bin/python
# -*-coding:utf-8-*-
import os
import numpy as np
import pandas as pd

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

# from dir_info import ylabel_data_dir
from dir_info import label_data_dir
from dir_info import cache_data_dir

from base.KeyTransform import stock_code_map
from base.KeyTransform import trade_date_map
from base.KeyTransform import ci1_code_map


def main(db = 'validation'):
    data_reader = DataReader(db=db)

    ##
    processed_trade_date_data = data_reader.read_basic_data_table('processed_trade_date_data')

    processed_trade_date_data = processed_trade_date_data.set_index('natural_date')

    processed_trade_date_data['label_date'] = processed_trade_date_data['trade_date'].copy()

    processed_trade_date_data.loc[processed_trade_date_data['trade_date']> processed_trade_date_data.index, 'label_date'] = \
        processed_trade_date_data.loc[processed_trade_date_data['trade_date']> processed_trade_date_data.index, 'last_trade_date']

    # TODO - 读取label日期
    target_construction_label = pd.read_hdf(os.path.join(cache_data_dir, 'X_financial_statement_data_target_construction_label.h5'))

    target_construction_label['label_date'] = \
        processed_trade_date_data.loc[target_construction_label.index.get_level_values('report_date'), 'label_date'].values

    target_construction_label = target_construction_label.reset_index()

    target_construction_label = trade_date_map(target_construction_label, key='report_date', reverse=False, is_financial_date=True)
    target_construction_label = trade_date_map(target_construction_label, key='label_date', reverse=False)
    target_construction_label = stock_code_map(target_construction_label, key='stock_code', reverse=False)

    target_construction_label = target_construction_label.set_index(['label_date', 'stock_code'])

    ## TODO - load label data
    concat_label_data = pd.read_hdf(os.path.join(label_data_dir, 'label_data.h5'))

    target_construction_label = pd.concat([target_construction_label, concat_label_data],
                                          axis=1,
                                          join_axes=[target_construction_label.index])

    target_construction_label = target_construction_label.reset_index()

    target_construction_label = target_construction_label.set_index(['report_date', 'stock_code'])

    print(target_construction_label.tail())

    ## TODO - SAVE
    save_filename = 'financial_label_data'

    target_construction_label.to_hdf(os.path.join(label_data_dir, save_filename + '.h5'),
                                     key=save_filename)


if __name__ == "__main__":
    main()


